"""
双流体模型的基本方程和数值求解方法

该模块实现了基于两流体模型的：
1. 基本守恒方程（质量、动量、能量）
2. 数值离散化方法
3. Picard-Krylov求解器
"""

import numpy as np
from typing import Dict, Tuple, List, Optional
from dataclasses import dataclass
from scipy.sparse import diags, csr_matrix
from scipy.sparse.linalg import spsolve
from .pgmres_solver import PGMRESSolver, PhysicalPreconditioner
from .fluid_properties import FluidProperties
from .closure_relations import ClosureRelations

@dataclass
class FluidProperties:
    """流体物性参数"""
    density: float  # 密度
    viscosity: float  # 粘度
    conductivity: float  # 导热系数
    specific_heat: float  # 比热容
    enthalpy: float  # 焓

class TwoFluidModel:
    def __init__(self, config: Dict, fluid_props: FluidProperties):
        """
        初始化双流体模型
        
        参数:
            config: 配置字典，包含几何和计算参数
            fluid_props: 流体物性计算器
        """
        # 保存配置和物性计算器
        self.config = config
        self.fluid_props = fluid_props
        
        # 网格参数
        self.n_axial = config['n_axial']  # 轴向网格数
        self.n_channels = config['n_channels']  # 子通道数
        
        # 初始化场变量 - 液相
        self.alpha_l = np.ones((self.n_axial, self.n_channels))  # 液相分数
        self.u_l = np.zeros((self.n_axial, self.n_channels))  # 液相轴向速度
        self.v_l = np.zeros((self.n_axial, self.n_channels))  # 液相横向速度
        self.h_l = np.zeros((self.n_axial, self.n_channels))  # 液相焓
        self.temp_l = np.zeros((self.n_axial, self.n_channels))  # 液相温度
        self.rho_l = np.zeros((self.n_axial, self.n_channels))  # 液相密度
        
        # 初始化场变量 - 气相
        self.alpha_g = np.zeros((self.n_axial, self.n_channels))  # 气相分数
        self.u_g = np.zeros((self.n_axial, self.n_channels))  # 气相轴向速度
        self.v_g = np.zeros((self.n_axial, self.n_channels))  # 气相横向速度
        self.h_g = np.zeros((self.n_axial, self.n_channels))  # 气相焓
        self.temp_g = np.zeros((self.n_axial, self.n_channels))  # 气相温度
        self.rho_g = np.zeros((self.n_axial, self.n_channels))  # 气相密度
        
        # 压力场
        self.pressure = np.ones((self.n_axial, self.n_channels)) * config['p_inlet']  # 压力场
        
        # 温度场
        self.temp_wall = np.zeros((self.n_axial, self.n_channels))  # 壁面温度
        self.temp_sat = np.zeros((self.n_axial, self.n_channels))  # 饱和温度
        
        # 热流和相变
        self.q_wall = np.zeros((self.n_axial, self.n_channels))  # 壁面热流密度
        self.gamma = np.zeros((self.n_axial, self.n_channels))  # 相变质量流量
        
        # 摩擦力
        self.f_wall_l = np.zeros((self.n_axial, self.n_channels))  # 液相壁面摩擦力
        self.f_wall_g = np.zeros((self.n_axial, self.n_channels))  # 气相壁面摩擦力
        self.f_i = np.zeros((self.n_axial, self.n_channels))  # 界面剪切力
        
        # 几何参数
        self.hydraulic_diameter = config['hydraulic_diameter']  # 水力直径
        
        # 初始化封闭关系计算器
        self.closure = ClosureRelations(config)
        
        # 存储上一次迭代的结果
        self._prev_alpha = np.copy(self.alpha_l)
        self._prev_u = np.copy(self.u_l)
        self._prev_v = np.copy(self.v_l)
        self._prev_pressure = np.copy(self.pressure)
        self._prev_temp = np.copy(self.temp_l)
        self._prev_h = np.copy(self.h_l)
        self._prev_rho = self.get_density('l', 0, 0)  # 示例索引
        
    def mass_conservation(self, k: str, i: int, j: int, dt: float) -> float:
        """
        计算质量守恒方程
        
        参数:
            k: 相态 ('l' 或 'g')
            i: 轴向索引
            j: 通道索引
            dt: 时间步长
            
        返回:
            质量守恒残差
        """
        # 获取相关物理量
        alpha = self.alpha_l if k == 'l' else self.alpha_g
        u = self.u_l if k == 'l' else self.u_g
        v = self.v_l if k == 'l' else self.v_g
        rho = self.get_density(k, i, j)
        
        # 时间导数项
        d_mass_dt = (alpha[i,j] * rho - self._prev_alpha[i,j] * self._prev_rho) / dt
        
        # 轴向对流项
        d_mass_dx = self.axial_mass_flux_derivative(k, i, j)
        
        # 横向对流项
        d_mass_dy = self.lateral_mass_flux_derivative(k, i, j)
        
        # 相变源项
        gamma = self.phase_change_rate(k, i, j)
        
        # 湍流混合项
        t_mk = self.turbulent_mixing_term(k, i, j)
        
        return d_mass_dt + d_mass_dx + d_mass_dy - gamma - t_mk
        
    def axial_momentum(self, k: str, i: int, j: int, dt: float) -> float:
        """
        计算轴向动量守恒方程
        
        参数:
            k: 相态 ('l' 或 'g')
            i: 轴向索引
            j: 通道索引
            dt: 时间步长
            
        返回:
            轴向动量守恒残差
        """
        # 获取相关物理量
        alpha = self.alpha_l if k == 'l' else self.alpha_g
        u = self.u_l if k == 'l' else self.u_g
        rho = self.get_density(k, i, j)
        
        # 时间导数项
        d_momentum_dt = (alpha[i,j] * rho * u[i,j] - 
                        self._prev_alpha[i,j] * self._prev_rho * self._prev_u[i,j]) / dt
        
        # 对流项
        d_momentum_dx = self.axial_momentum_flux_derivative(k, i, j)
        
        # 压力梯度项
        dp_dx = (self.pressure[i+1,j] - self.pressure[i,j]) / self.dx
        pressure_term = -alpha[i,j] * dp_dx
        
        # 重力项
        gravity_term = -alpha[i,j] * rho * 9.81
        
        # 壁面摩擦力
        f_wall = self.wall_friction(k, i, j)
        
        # 界面剪切力
        f_interface = self.interface_shear(k, i, j)
        
        return (d_momentum_dt + d_momentum_dx + pressure_term + 
                gravity_term + f_wall + f_interface)
        
    def lateral_momentum(self, k: str, i: int, j: int, dt: float) -> float:
        """
        计算横向动量守恒方程
        
        参数:
            k: 相态 ('l' 或 'g')
            i: 轴向索引
            j: 通道索引
            dt: 时间步长
            
        返回:
            横向动量守恒残差
        """
        # 获取相关物理量
        alpha = self.alpha_l if k == 'l' else self.alpha_g
        v = self.v_l if k == 'l' else self.v_g
        rho = self.get_density(k, i, j)
        
        # 时间导数项
        d_momentum_dt = (alpha[i,j] * rho * v[i,j] - 
                        self._prev_alpha[i,j] * self._prev_rho * self._prev_v[i,j]) / dt
        
        # 对流项（简化处理，忽略正交方向动量）
        d_momentum_dy = self.lateral_momentum_flux_derivative(k, i, j)
        
        # 压力梯度项
        dp_dy = self.lateral_pressure_gradient(i, j)
        pressure_term = -alpha[i,j] * dp_dy
        
        # 壁面摩擦力
        f_wall = self.lateral_wall_friction(k, i, j)
        
        # 界面剪切力
        f_interface = self.lateral_interface_shear(k, i, j)
        
        # 湍流混合动量源项
        t_mk = self.turbulent_mixing_momentum(k, i, j)
        
        return (d_momentum_dt + d_momentum_dy + pressure_term + 
                f_wall + f_interface + t_mk)
        
    def energy_conservation(self, k: str, i: int, j: int, dt: float) -> float:
        """
        计算能量守恒方程
        
        参数:
            k: 相态 ('l' 或 'g')
            i: 轴向索引
            j: 通道索引
            dt: 时间步长
            
        返回:
            能量守恒残差
        """
        # 获取相关物理量
        alpha = self.alpha_l if k == 'l' else self.alpha_g
        h = self.h_l if k == 'l' else self.h_g
        u = self.u_l if k == 'l' else self.u_g
        rho = self.get_density(k, i, j)
        
        # 时间导数项
        d_energy_dt = (alpha[i,j] * rho * h[i,j] - 
                      self._prev_alpha[i,j] * self._prev_rho * self._prev_h[i,j]) / dt
        
        # 轴向对流项
        d_energy_dx = self.axial_energy_flux_derivative(k, i, j)
        
        # 横向对流项
        d_energy_dy = self.lateral_energy_flux_derivative(k, i, j)
        
        # 壁面传热项
        q_wall = self._calculate_wall_heat_flux() if k == 'l' else 0
        
        # 相变能量项
        gamma = self.phase_change_rate(k, i, j)
        h_sat = self._get_saturation_enthalpy(self.pressure[i,j])
        phase_change_energy = gamma * h_sat
        
        # 湍流混合能量项
        t_mk = self.turbulent_mixing_term(k, i, j)
        
        return d_energy_dt + d_energy_dx + d_energy_dy - q_wall - phase_change_energy - t_mk

    def axial_mass_flux_derivative(self, k: str, i: int, j: int) -> float:
        """计算轴向质量通量导数"""
        alpha = self.alpha_l if k == 'l' else self.alpha_g
        u = self.u_l if k == 'l' else self.u_g
        rho = self.get_density(k, i, j)
        
        if i > 0:
            flux_prev = alpha[i-1,j] * rho * u[i-1,j]
            flux_curr = alpha[i,j] * rho * u[i,j]
            return (flux_curr - flux_prev) / self.dx
        return 0

    def lateral_mass_flux_derivative(self, k: str, i: int, j: int) -> float:
        """计算横向质量通量导数"""
        alpha = self.alpha_l if k == 'l' else self.alpha_g
        v = self.v_l if k == 'l' else self.v_g
        rho = self.get_density(k, i, j)
        
        if j > 0:
            flux_prev = alpha[i,j-1] * rho * v[i,j-1]
            flux_curr = alpha[i,j] * rho * v[i,j]
            return (flux_curr - flux_prev) / self.dy
        return 0

    def phase_change_rate(self, k: str, i: int, j: int) -> float:
        """计算相变率"""
        # 获取当前温度和饱和温度
        T = self.temp_l[i,j] if k == 'l' else self.temp_g[i,j]
        T_sat = self.temp_sat[i,j]
        
        # 获取焓值
        h = self.h_l[i,j] if k == 'l' else self.h_g[i,j]
        h_sat = self._get_saturation_enthalpy(self.pressure[i,j])
        
        # 计算相变率
        if k == 'l':
            if T > T_sat:  # 过热液体
                return self.closure._superheated_liquid_mass_flux(h, h_sat, self.r, self.cp)
            else:  # 过冷液体
                return self.closure._subcooled_liquid_mass_flux(h, h_sat, self.r, self.cp)
        else:
            if T < T_sat:  # 过冷蒸汽
                return self.closure._subcooled_vapor_mass_flux(h, h_sat, self.r, self.cp)
            else:  # 过热蒸汽
                return self.closure._superheated_vapor_mass_flux(h, h_sat, self.r, self.cp)

    def turbulent_mixing_term(self, k: str, i: int, j: int) -> float:
        """计算湍流混合项"""
        return self.closure.turbulent_mixing_momentum(k, i, j)

    def _get_saturation_enthalpy(self, pressure: float) -> float:
        """获取饱和焓值"""
        return self.fluid_props.get_saturation_enthalpy(pressure)

    def axial_energy_flux_derivative(self, k: str, i: int, j: int) -> float:
        """计算轴向能量通量导数"""
        alpha = self.alpha_l if k == 'l' else self.alpha_g
        u = self.u_l if k == 'l' else self.u_g
        h = self.h_l if k == 'l' else self.h_g
        rho = self.get_density(k, i, j)
        
        if i > 0:
            flux_prev = alpha[i-1,j] * rho * u[i-1,j] * h[i-1,j]
            flux_curr = alpha[i,j] * rho * u[i,j] * h[i,j]
            return (flux_curr - flux_prev) / self.dx
        return 0

    def lateral_energy_flux_derivative(self, k: str, i: int, j: int) -> float:
        """计算横向能量通量导数"""
        alpha = self.alpha_l if k == 'l' else self.alpha_g
        v = self.v_l if k == 'l' else self.v_g
        h = self.h_l if k == 'l' else self.h_g
        rho = self.get_density(k, i, j)
        
        if j > 0:
            flux_prev = alpha[i,j-1] * rho * v[i,j-1] * h[i,j-1]
            flux_curr = alpha[i,j] * rho * v[i,j] * h[i,j]
            return (flux_curr - flux_prev) / self.dy
        return 0

    def solve_picard_krylov(self, dt: float, max_iter: int = 20, 
                           tol: float = 1e-2) -> bool:
        """
        使用Picard-Krylov方法求解方程组
        
        参数:
            dt: 时间步长
            max_iter: 最大迭代次数
            tol: 收敛容差
            
        返回:
            是否收敛
        """
        for iter in range(max_iter):
            # 保存上一次迭代的结果
            self._store_previous_iteration()
            
            # 1. 计算加热棒与流场间的热流
            self._calculate_wall_heat_flux()
            
            # 2. 求解动量方程获得质量流率
            self._solve_momentum_equations()
            
            # 3. 使用PGMRES方法求解质量-能量方程
            self._solve_mass_energy_equations()
            
            # 4. 计算两相本构方程
            self._update_constitutive_relations()
            
            # 检查收敛性
            if self._check_convergence(tol):
                return True
                
            # 如果迭代次数达到最大值但仍未收敛
            if iter == max_iter - 1:
                return False
                
    def _store_previous_iteration(self):
        """存储上一次迭代的结果"""
        self._prev_alpha = np.copy(self.alpha_l)
        self._prev_u = np.copy(self.u_l)
        self._prev_v = np.copy(self.v_l)
        self._prev_pressure = np.copy(self.pressure)
        self._prev_temp = np.copy(self.temp_l)
        self._prev_h = np.copy(self.h_l)
        self._prev_rho = self.get_density('l', 0, 0)  # 示例索引
        
    def _calculate_wall_heat_flux(self):
        """计算壁面热流"""
        for i in range(self.n_axial):
            for j in range(self.n_channels):
                # 获取当前子通道的几何参数
                channel_geom = self.geometry['channels'][j]
                
                # 计算传热系数
                h_v = self.closure.heat_transfer_coefficient(
                    'vapor', self.pressure[i,j], self.temp_wall[i,j],
                    self.temp_v[i,j], self.temp_sat[i,j])
                h_l = self.closure.heat_transfer_coefficient(
                    'liquid', self.pressure[i,j], self.temp_wall[i,j],
                    self.temp_l[i,j], self.temp_sat[i,j])
                
                # 更新热流密度
                self.q_wall[i,j] = h_v * (self.temp_wall[i,j] - self.temp_v[i,j]) + \
                                  h_l * (self.temp_wall[i,j] - self.temp_l[i,j])
                
    def _solve_momentum_equations(self):
        """求解动量方程"""
        for i in range(self.n_axial):
            for j in range(self.n_channels):
                # 求解轴向动量方程
                self.u_l[i,j] = self._solve_axial_momentum('l', i, j)
                self.u_g[i,j] = self._solve_axial_momentum('g', i, j)
                
                # 求解横向动量方程
                self.v_l[i,j] = self._solve_lateral_momentum('l', i, j)
                self.v_g[i,j] = self._solve_lateral_momentum('g', i, j)
                
    def _solve_mass_energy_equations(self):
        """使用PGMRES方法求解质量-能量方程组"""
        # 构建雅可比矩阵和残差向量
        jacobian = self._build_jacobian_matrix()
        residual = self._build_residual_vector()
        
        # 初始化PGMRES求解器
        preconditioner = PhysicalPreconditioner(self.n_axial, self.n_channels)
        solver = PGMRESSolver(preconditioner)
        
        # 求解方程组
        delta_x, converged = solver.solve(jacobian, residual)
        
        if not converged:
            print("警告: PGMRES求解未收敛")
            
        # 更新变量
        self._update_variables(delta_x)
        
    def _update_constitutive_relations(self):
        """更新本构关系"""
        for i in range(self.n_axial):
            for j in range(self.n_channels):
                # 更新相变质量流量
                self.gamma[i,j] = self.closure.phase_change_mass_flux(
                    self.pressure[i,j], self.temp_l[i,j], self.temp_v[i,j],
                    self.h_l[i,j], self.h_v[i,j])
                
                # 更新壁面摩擦力
                self.f_wall_l[i,j], self.f_wall_g[i,j] = \
                    self.closure.wall_interface_friction(
                        'liquid', self.u_l[i,j], self.hydraulic_diameter,
                        self.rho_l[i,j], self.mu_l[i,j])
                
                # 更新界面剪切力
                self.f_i[i,j] = self.closure.interface_shear(
                    self.u_l[i,j], self.u_g[i,j], self.alpha_l[i,j],
                    self.rho_l[i,j], self.rho_g[i,j])
                
    def _check_convergence(self, tol: float) -> bool:
        """
        检查迭代收敛性
        
        参数:
            tol: 收敛容差
            
        返回:
            是否收敛
        """
        # 计算主要变量的相对变化
        delta_p = np.max(np.abs((self.pressure - self._prev_pressure) / 
                               self._prev_pressure))
        delta_alpha = np.max(np.abs((self.alpha_l - self._prev_alpha) / 
                                   self._prev_alpha))
        delta_h = np.max(np.abs((self.h_l - self._prev_h) / self._prev_h))
        
        # 检查所有变量是否都满足收敛条件
        return max(delta_p, delta_alpha, delta_h) < tol
        
    def _build_jacobian_matrix(self) -> csr_matrix:
        """
        构建雅可比矩阵
        
        返回:
            雅可比矩阵(CSR格式)
        """
        n = self.n_axial * self.n_channels * 4  # 4个变量(p, α, h_l, h_g)
        
        # 使用有限差分法计算雅可比矩阵
        eps = 1e-6  # 扰动大小
        
        # 初始化稀疏矩阵的数据
        data = []
        rows = []
        cols = []
        
        for i in range(n):
            # 确定当前变量类型和位置
            var_type = i % 4
            channel = (i // 4) % self.n_channels
            axial = (i // (4 * self.n_channels))
            
            # 计算当前方程对所有变量的偏导数
            for j in range(n):
                # 使用前向差分计算偏导数
                delta = np.zeros(n)
                delta[j] = eps
                
                # 计算扰动后的残差
                res1 = self._calculate_residual()
                self._apply_perturbation(delta)
                res2 = self._calculate_residual()
                self._remove_perturbation(delta)
                
                # 计算偏导数
                deriv = (res2[i] - res1[i]) / eps
                
                # 如果偏导数足够大，添加到稀疏矩阵
                if abs(deriv) > 1e-10:
                    data.append(deriv)
                    rows.append(i)
                    cols.append(j)
                    
        return csr_matrix((data, (rows, cols)), shape=(n, n))
        
    def _build_residual_vector(self) -> np.ndarray:
        """
        构建残差向量
        
        返回:
            残差向量
        """
        return self._calculate_residual()
        
    def _update_variables(self, delta_x: np.ndarray):
        """
        更新变量
        
        参数:
            delta_x: 变量的更新量
        """
        n_vars = 4  # 每个节点4个变量
        
        for i in range(self.n_axial):
            for j in range(self.n_channels):
                idx = (i * self.n_channels + j) * n_vars
                
                # 更新压力
                self.pressure[i,j] += delta_x[idx]
                
                # 更新空泡分数
                self.alpha_l[i,j] += delta_x[idx + 1]
                self.alpha_g[i,j] = 1 - self.alpha_l[i,j]
                
                # 更新焓
                self.h_l[i,j] += delta_x[idx + 2]
                self.h_g[i,j] += delta_x[idx + 3]
                
                # 更新温度(使用状态方程)
                self.temp_l[i,j] = self.fluid_props.get_temperature(
                    self.pressure[i,j], self.h_l[i,j], 'liquid')
                self.temp_g[i,j] = self.fluid_props.get_temperature(
                    self.pressure[i,j], self.h_g[i,j], 'vapor')
                
    def _calculate_residual(self) -> np.ndarray:
        """
        计算方程组的残差
        
        返回:
            残差向量
        """
        n = self.n_axial * self.n_channels * 4
        residual = np.zeros(n)
        
        for i in range(self.n_axial):
            for j in range(self.n_channels):
                idx = (i * self.n_channels + j) * 4
                
                # 质量守恒残差
                residual[idx] = self.mass_conservation('l', i, j, self.dt)
                
                # 空泡分数残差
                residual[idx + 1] = self.alpha_l[i,j] + self.alpha_g[i,j] - 1
                
                # 能量守恒残差
                residual[idx + 2] = self.energy_conservation('l', i, j, self.dt)
                residual[idx + 3] = self.energy_conservation('g', i, j, self.dt)
                
        return residual
        
    def _apply_perturbation(self, delta: np.ndarray):
        """应用扰动"""
        n_vars = 4
        for i in range(self.n_axial):
            for j in range(self.n_channels):
                idx = (i * self.n_channels + j) * n_vars
                self.pressure[i,j] += delta[idx]
                self.alpha_l[i,j] += delta[idx + 1]
                self.h_l[i,j] += delta[idx + 2]
                self.h_g[i,j] += delta[idx + 3]
                
    def _remove_perturbation(self, delta: np.ndarray):
        """移除扰动"""
        n_vars = 4
        for i in range(self.n_axial):
            for j in range(self.n_channels):
                idx = (i * self.n_channels + j) * n_vars
                self.pressure[i,j] -= delta[idx]
                self.alpha_l[i,j] -= delta[idx + 1]
                self.h_l[i,j] -= delta[idx + 2]
                self.h_g[i,j] -= delta[idx + 3]
        
    def _solve_axial_momentum(self, k: str, i: int, j: int) -> float:
        """
        求解轴向动量方程
        
        参数:
            k: 相态 ('l' 或 'g')
            i: 轴向索引
            j: 通道索引
            
        返回:
            轴向速度
        """
        # 获取相关物理量
        alpha = self.alpha_l if k == 'l' else self.alpha_g
        rho = self.get_density(k, i, j)
        u = self.u_l if k == 'l' else self.u_g
        v = self.v_l if k == 'l' else self.v_g
        
        # 时间导数项
        d_momentum_dt = (alpha[i,j] * rho * u[i,j] - 
                        self._prev_alpha[i,j] * self._prev_rho * self._prev_u[i,j]) / self.dt
        
        # 对流项
        # 轴向对流
        if i > 0:
            d_momentum_dx = (alpha[i,j] * rho * u[i,j] * u[i,j] - 
                           alpha[i-1,j] * rho * u[i-1,j] * u[i-1,j]) / self.dx
        else:
            d_momentum_dx = 0
            
        # 横向对流
        if j > 0:
            d_momentum_dy = (alpha[i,j] * rho * u[i,j] * v[i,j] - 
                           alpha[i,j-1] * rho * u[i,j-1] * v[i,j-1]) / self.dy
        else:
            d_momentum_dy = 0
            
        # 压力梯度项
        if i < self.n_axial - 1:
            dp_dx = (self.pressure[i+1,j] - self.pressure[i,j]) / self.dx
        else:
            dp_dx = 0
            
        # 重力项
        gravity = -alpha[i,j] * rho * 9.81 if k == 'l' else -alpha[i,j] * rho * 9.81
        
        # 壁面摩擦力
        f_wall = self.f_wall_l[i,j] if k == 'l' else self.f_wall_g[i,j]
        
        # 界面剪切力
        f_i = self.f_i[i,j] if k == 'l' else -self.f_i[i,j]
        
        # 相变动量项
        gamma = self.gamma[i,j] if k == 'l' else -self.gamma[i,j]
        
        # 计算新的速度
        u_new = u[i,j] - (d_momentum_dt + d_momentum_dx + d_momentum_dy + 
                         alpha[i,j] * dp_dx + gravity + f_wall - f_i - gamma * u[i,j]) * self.dt
        
        return u_new
        
    def _solve_lateral_momentum(self, k: str, i: int, j: int) -> float:
        """
        求解横向动量方程
        
        参数:
            k: 相态 ('l' 或 'g')
            i: 轴向索引
            j: 通道索引
            
        返回:
            横向速度
        """
        # 获取相关物理量
        alpha = self.alpha_l if k == 'l' else self.alpha_g
        rho = self.get_density(k, i, j)
        u = self.u_l if k == 'l' else self.u_g
        v = self.v_l if k == 'l' else self.v_g
        
        # 时间导数项
        d_momentum_dt = (alpha[i,j] * rho * v[i,j] - 
                        self._prev_alpha[i,j] * self._prev_rho * self._prev_v[i,j]) / self.dt
        
        # 对流项 - 根据子通道分析方法简化，只考虑轴向对流
        if i > 0:
            d_momentum_dx = (alpha[i,j] * rho * v[i,j] * u[i,j] - 
                           alpha[i-1,j] * rho * v[i-1,j] * u[i-1,j]) / self.dx
        else:
            d_momentum_dx = 0
            
        # 压力梯度项
        if j < self.n_channels - 1:
            dp_dy = (self.pressure[i,j+1] - self.pressure[i,j]) / self.dy
        else:
            dp_dy = 0
            
        # 壁面摩擦力
        f_wall = self.f_wall_l[i,j] if k == 'l' else self.f_wall_g[i,j]
        
        # 界面剪切力
        f_i = self.f_i[i,j] if k == 'l' else -self.f_i[i,j]
        
        # 相变动量项
        gamma = self.gamma[i,j] if k == 'l' else -self.gamma[i,j]
        
        # 湍流混合动量项
        t_mix = self.closure.turbulent_mixing_momentum(k, i, j)
        
        # 计算新的速度
        v_new = v[i,j] - (d_momentum_dt + d_momentum_dx + alpha[i,j] * dp_dy + 
                         f_wall - f_i - gamma * v[i,j] + t_mix) * self.dt
        
        return v_new
        
    def get_density(self, k: str, i: int, j: int) -> float:
        """
        计算指定相态和位置的密度
        
        参数:
            k: 相态 ('l' 或 'g')
            i: 轴向索引
            j: 通道索引
            
        返回:
            密度 (kg/m³)
        """
        if k == 'l':
            # 液相密度
            h = self.h_l[i,j]
            quality = self.fluid_props.get_quality(self.pressure[i,j], h)
            if quality == 0:
                # 单相液体
                return self.fluid_props.get_liquid_density(self.pressure[i,j], self.temp_l[i,j])
            else:
                # 两相流中的液相
                return self.fluid_props.get_liquid_density(self.pressure[i,j], self.temp_sat[i,j])
        else:
            # 气相密度
            h = self.h_g[i,j]
            quality = self.fluid_props.get_quality(self.pressure[i,j], h)
            if quality == 1:
                # 单相气体
                return self.fluid_props.get_vapor_density(self.pressure[i,j], self.temp_g[i,j])
            else:
                # 两相流中的气相
                return self.fluid_props.get_vapor_density(self.pressure[i,j], self.temp_sat[i,j])